from models import Status
from utils import Neo4jUtil
import pandas as pd
from question.kb_service.milvus_service import MilvusService
from config.config_env import KB_NAME

# 知识库
kb_service = MilvusService(KB_NAME)


# 审查文件
def examine_file(df):
    columns = df.columns.tolist()
    columns_string = ''.join(columns)
    if columns_string == '设备部件检查内容故障现象排故流程故障原因维修方案':
        return True
    else:
        return False


class File2DatabaseUtil(object):

    def __init__(self):
        self.g = Neo4jUtil()
        # 主数据库
        self.db = self.g.graph
        # 从数据库
        self.synchronize_db = self.g.from_graph
        self.triples = [['设备', '包含', '部件'], ['部件', '产生', '故障现象'], ['部件', '包括', '检查内容'],
                        ['检查内容', '涉及', '故障现象']
            , ['故障原因', '导致', '故障现象'], ['故障原因', '相关', '部件'], ['故障原因', '建议', '维修方案'],
                        ['故障现象', '排故', '排故流程'], ['排故流程', '对应', '故障原因']]

    # 数据入库，要求与模版匹配
    def import_file(self, file_path, id, current_app):
        df = pd.read_excel(file_path)
        if examine_file(df):
            for triple in self.triples:
                for index, row in df.iterrows():
                    cql = '''MERGE (a:%s {name: '%s'})
                            MERGE (b:%s {name: '%s'})
                            MERGE (a)-[r:%s]-> (b)
                            ''' % (triple[0], row[triple[0]], triple[2], row[triple[2]], triple[1])
                    # 主数据库录入
                    self.db.run(cql)
                    # 从数据库录入
                    self.synchronize_db.run(cql)
            # 文件入知识库
            kb_service.insert_doc_to_vector(file_path)
            with current_app.app_context():
                s = Status.query.get(id)
                s.change_status()
            return 1
        else:
            return 0
